import cv2
import os
import os.path
import xml.etree.cElementTree as ET


def draw(image_path, xml_path, root_saved_path):
    "图片根据标注画框"
    src_img_path = image_path
    src_ann_path = xml_path
    for file in os.listdir(src_ann_path):
        file_name, suffix = os.path.splitext(file)
        print(suffix)
        if suffix == '.xml':
            print(file)
            xml_path = os.path.join(src_ann_path, file)
            image_path = os.path.join(src_img_path, file_name + '.jpg')
            img = cv2.imread(image_path)
            tree = ET.parse(xml_path)
            root = tree.getroot()
            for obj in root.iter('object'):
                name = obj.find('name').text
                xml_box = obj.find('bndbox')
                x1 = int(float(xml_box.find('xmin').text))
                x2 = int(float(xml_box.find('xmax').text))
                y1 = int(float(xml_box.find('ymin').text))
                y2 = int(float(xml_box.find('ymax').text))
                cv2.rectangle(img, (x1, y1), (x2, y2), (255, 0, 0), thickness=2)
                cv2.putText(img, name, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.75, (0, 0, 255), thickness=2)
            cv2.imwrite(os.path.join(root_saved_path, file_name + '.jpg'), img)


def find(image_path, xml_path):
    "图片根据标注画框"
    src_img_path = image_path
    src_ann_path = xml_path
    i = 0
    for file in os.listdir(src_ann_path):
        # file是xml文件名
        file_name, suffix = os.path.splitext(file)

        xml_path = os.path.join(src_ann_path, file)
        tree = ET.parse(xml_path)
        root = tree.getroot()
        filename = root.find('filename').text  # jpg
        filename = filename[:-4]

        filename = filename[-8:]
        file_name = file_name[-8:]
        if (filename != file_name):
            i += 1
            print('-' * 20 + f'xml名字：{filename}\n文件名：{file_name}')
    print(f'总共有{i}个xml文件名与内部filename不同')


if __name__ == '__main__':
    image_path = '../images'
    xml_path = '../image_annotation'
    root_save_path = '../images2'
    # 通过xml绘制带框框的图片
    draw(image_path, xml_path, root_save_path)
    # 检查xml名称是否与内部filename相同
    find(image_path, xml_path)
